"""
训练模型
"""
import time

from parl.algorithms import DDQN

from Examples.DQNExamples.Dimension2D.Dim2DAgent import Dim2DAgent
from Examples.DQNExamples.HollowKnight2D.HK2DModel import HK2DModel
from Examples.DQNExamples.HollowKnight2D.HK2DRPM import HK2DRPM
from Examples.DQNExamples.HollowKnight2D.HKEnv2D import HKEnv2D
from Examples.DQNExamples.HollowKnight2D.HKTrain2D import HKTraining2D

MEMORY_SIZE = 200  # 记忆库大小
BATCH_SIZE = 32  # 批大小
GAMMA = 0.999  # 奖励衰减率
LEARNING_RATE = 0.001  # 学Z习率
E_GREED = 1  # 贪心率
E_GREED_DECREMENT = 1e-6  # 贪心率衰减率
UPDATE_TARGET_STEPS = 200  # 更新目标网络的步数
MAX_EPISODES = 90  # 最大回合数
MAX_TURN = 1000  # 最大轮数
MAX_STEPS = 200000  # 训练每回合最大步数, 避免训练时间过长
MAX_TEST_STEPS = 20000  # 测试每回合最大步数, 避免时间过长
MEMORY_WARMUP_SIZE = 64  # replay memory预热大小, 需要预存一些经验数据后才开始训练
TEST_FREQ = 100  # 测试的频率, 每隔多少回合测试一次
FRAME_SKIPPING = 10  # 每隔多少帧采样一次
TARGET_UPDATE_FREQ = 5  # 每隔多少次训练更新一次target_net

for turn in range(MAX_TURN):
    # 创建环境
    env = HKEnv2D()

    # 获取动作维度和状态维度
    act_dim = 11
    move_dim = 7
    # obs_dim = env.observation_space.n
    obs_dim = (311, 311)

    # 经验回放池
    start = time.time()
    rmp_act = HK2DRPM(MEMORY_SIZE, obs_dim, 4)
    rmp_move = HK2DRPM(MEMORY_SIZE, obs_dim, 4)
    end = time.time()
    print("经验回放池创建耗时: ", end - start)

    # 创建模型
    start = time.time()
    model_act = HK2DModel(act_dim=act_dim, obs_dim=obs_dim)
    model_move = HK2DModel(act_dim=move_dim, obs_dim=obs_dim)
    end = time.time()
    print("模型创建耗时: ", end - start)

    # 创建算法
    start = time.time()
    algorithm_act = DDQN(model_act, gamma=GAMMA, lr=LEARNING_RATE)
    algorithm_move = DDQN(model_move, gamma=GAMMA, lr=LEARNING_RATE)
    end = time.time()
    print("算法创建耗时: ", end - start)

    # 创建Agent
    start = time.time()
    agent_act = Dim2DAgent(algorithm_act, obs_dim, act_dim,
                           e_greed=E_GREED,
                           e_greed_decrement=E_GREED_DECREMENT,
                           update_target_steps=UPDATE_TARGET_STEPS)
    agent_move = Dim2DAgent(algorithm_move, obs_dim, move_dim,
                            e_greed=E_GREED,
                            e_greed_decrement=E_GREED_DECREMENT,
                            update_target_steps=UPDATE_TARGET_STEPS)
    end = time.time()
    print("Agent创建耗时: ", end - start)

    # 创建训练
    train = HKTraining2D()

    # 开始训练
    train.test(
        agent_act=agent_act,
        agent_move=agent_move,
        env=env,
        rpm_move=rmp_move,
        rpm_act=rmp_act,
        max_steps=MAX_STEPS,
        act_path='model/big_fly_act.ckpt',
        move_path='model/big_fly_move.ckpt',
    )
